import os
import time
import cv2
import numpy as np
from ultralytics import YOLO
from picamera2 import Picamera2
import logging
import serial

# 配置日志
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('detection.log')
    ]
)
logger = logging.getLogger(__name__)

# 目标类别名称
TARGET_CLASS = "target_stripe"
CUSTOM_CLASS_NAMES = ['target_stripe']

# 串口配置
SERIAL_PORT = '/dev/ttyAMA0'
BAUD_RATE = 9600
ser = None

# 坐标系设置
SCREEN_ORIGIN = "BOTTOM_LEFT"
SCREEN_WIDTH = 640
SCREEN_HEIGHT = 480

# 坐标平滑参数
SMOOTHING_FACTOR = 0.2  # 值越小越平滑 (0.1-0.5)
smoothed_x = SCREEN_WIDTH // 2
smoothed_y = SCREEN_HEIGHT // 2

def release_camera_resources():
    """释放摄像头资源"""
    logger.info("释放摄像头资源...")
    os.system("sudo pkill -f libcamera")
    time.sleep(1)

def init_serial():
    """初始化串口连接"""
    global ser
    try:
        ser = serial.Serial(SERIAL_PORT, BAUD_RATE, timeout=1)
        logger.info(f"串口连接成功: {SERIAL_PORT} @ {BAUD_RATE}")
        return True
    except Exception as e:
        logger.error(f"串口连接失败: {str(e)}")
        return False

def calculate_coordinates(box_xyxy, width, height):
    """计算对角线交点坐标（左下角为原点）"""
    x1, y1, x2, y2 = box_xyxy
    center_x = (x1 + x2) // 2
    center_y = (y1 + y2) // 2
    
    if SCREEN_ORIGIN == "BOTTOM_LEFT":
        center_y = height - center_y
    return center_x, center_y

def smooth_coordinates(x, y):
    """使用指数加权移动平均平滑坐标"""
    global smoothed_x, smoothed_y
    smoothed_x = SMOOTHING_FACTOR * x + (1 - SMOOTHING_FACTOR) * smoothed_x
    smoothed_y = SMOOTHING_FACTOR * y + (1 - SMOOTHING_FACTOR) * smoothed_y
    return int(smoothed_x), int(smoothed_y)

def create_data_packet(x, y):
    """
    创建符合协议的数据包
    格式: [包头(0xFF), x高8位, x低8位, y高8位, y低8位, 包尾(0xFE)]
    """

 # 确保坐标在有效范围内
    x = max(0, min(x, SCREEN_WIDTH - 1))
    y = max(0, min(y, SCREEN_HEIGHT - 1))
    
    # 将坐标分解为高8位和低8位
    x_high = (x >> 8) & 0xFF
    x_low = x & 0xFF
    y_high = (y >> 8) & 0xFF
    y_low = y & 0xFF
    
    # 构建数据包
    packet = bytes([0xFF, x_high, x_low, y_high, y_low, 0xFE])
    return packet

def main():
    global smoothed_x, smoothed_y
    
    logger.info("开始加载模型...")
    model = None
    pt_path = "/home/cai/yolo/best.pt"
    onnx_path = "/home/cai/yolo/best.onnx"
    
    if not os.path.exists(pt_path) and not os.path.exists(onnx_path):
        logger.error("模型文件未找到！")
        return
    
    try:
        model = YOLO(pt_path)
        logger.info("PyTorch模型加载成功")
    except:
        try:
            model = YOLO(onnx_path)
            logger.info("ONNX模型加载成功")
        except:
            logger.error("模型加载失败")
            return
    
    # 初始化串口
    serial_connected = init_serial()
    
     # 摄像头初始化
    release_camera_resources()
    picam2 = Picamera2()
    config = picam2.create_video_configuration(
        main={"size": (SCREEN_WIDTH, SCREEN_HEIGHT), "format": "RGB888"}
    )
    picam2.configure(config)
    picam2.start()
    logger.info(f"摄像头已启动 ({SCREEN_WIDTH}x{SCREEN_HEIGHT})")

    frame_count = 0
    start_time = time.time()
    last_target_save_time = time.time()
    last_coord_send_time = time.time()
    target_save_interval = 5
    coord_send_interval = 1.0  # 1秒发送一次坐标
    
    # 存储最新的坐标
    latest_coord = (0, 0)
    has_valid_coord = False  # 是否有有效坐标
    coord_expire_time = 2.0  # 坐标有效期（秒），超过此时间未更新则停止发送

    logger.info("开始检测循环...")
    try:
        while True:
            current_time = time.time()
            frame = picam2.capture_array()
            frame_count += 1
            frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
            
            # 模型推理
            results = model(frame, imgsz=SCREEN_WIDTH, verbose=False)
            
            detection_occurred = False
            
            if results[0].boxes:
                for box in results[0].boxes:
                    class_idx = int(box.cls)
                    class_name = model.names[class_idx]
                    if class_name != TARGET_CLASS:
                        continue
                    
                    xyxy = box.xyxy.cpu().numpy().squeeze().astype(int)
                    logger.debug(f"检测到目标: {class_name}, 坐标: {xyxy}")

 
                    # 计算坐标
                    center_x, center_y = calculate_coordinates(xyxy, SCREEN_WIDTH, SCREEN_HEIGHT)
                    
                    # 平滑坐标
                    smoothed_x, smoothed_y = smooth_coordinates(center_x, center_y)
                    logger.info(f"原始坐标: ({center_x}, {center_y}) 平滑后: ({smoothed_x}, {smoothed_y})")
                    
                    # 更新最新坐标和有效期
                    latest_coord = (smoothed_x, smoothed_y)
                    has_valid_coord = True
                    last_detection_time = current_time  # 记录最后一次检测时间
                    detection_occurred = True
                    
                     # 保存图片
                    if current_time - last_target_save_time > target_save_interval:
                        cv2.imwrite(f"target_{int(current_time)}.jpg", frame_bgr)
                        logger.info(f"保存目标图片: target_{int(current_time)}.jpg")
                        last_target_save_time = current_time
                    
                    # 绘制检测框和坐标
                    cv2.rectangle(frame_bgr, (xyxy[0], xyxy[1]), (xyxy[2], xyxy[3]), (0, 255, 0), 2)
                    # 显示原始坐标和平滑坐标
                    cv2.putText(frame_bgr, 
                               f"Raw: ({center_x}, {center_y})", 
                               (xyxy[0], xyxy[1]-30), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
                    cv2.putText(frame_bgr, 
                               f"Smooth: ({smoothed_x}, {smoothed_y})", 
                               (xyxy[0], xyxy[1]-10), 
                               cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)
            
            # 检查坐标是否过期
            if has_valid_coord and (current_time - last_detection_time) > coord_expire_time:
                logger.info("坐标已过期，停止发送")
                has_valid_coord = False
                # 重置平滑坐标到中心点
                smoothed_x = SCREEN_WIDTH // 2
                smoothed_y = SCREEN_HEIGHT // 2
            
            # 每秒发送一次坐标（只在有有效坐标时发送）
            if current_time - last_coord_send_time >= coord_send_interval:
                if serial_connected and ser.is_open and has_valid_coord:
                    data_packet = create_data_packet(*latest_coord)
                    ser.write(data_packet)
                    # 日志记录发送的原始字节
                    hex_packet = ' '.join([f'{b:02X}' for b in data_packet])
                    logger.info(f"发送坐标: ({latest_coord[0]}, {latest_coord[1]}) - 数据包: {hex_packet}")
                
                # 重置发送计时器
                last_coord_send_time = current_time
            
            # 调试：每30帧输出FPS
            if frame_count % 30 == 0:
                fps = frame_count / (current_time - start_time)
                logger.info(f"当前FPS: {fps:.1f}")

            time.sleep(0.03)

    except KeyboardInterrupt:
        logger.info("用户中断")
    except Exception as e:
        logger.error(f"运行时错误: {str(e)}", exc_info=True)
    finally:
        picam2.stop()
        if ser and ser.is_open:
            ser.close()
        logger.info("程序结束")

if __name__ == "__main__":
    main()
                    